#include "mapping/octomap_multi_scan.h"

Gen_grid_map::Gen_grid_map(ros::NodeHandle& nh_){

    nh = nh_;

    sub_cloud            = nh.subscribe<cloud_pose_msg::CloudPose>("/cloud_enu",10000,&Gen_grid_map::cloud_callback,this);
    pub_ground_cloud     = nh.advertise<sensor_msgs::PointCloud2>("ground_cloud",1000);
    pub_non_ground_cloud = nh.advertise<sensor_msgs::PointCloud2>("non_ground_cloud",1000);
    pub_full_cloud       = nh.advertise<sensor_msgs::PointCloud2>("full_cloud",1000);
    pub_grid_map         = nh.advertise<nav_msgs::OccupancyGrid>("grid_map",1000);

    std::string yaml_path = std::string(PROJECT_PATH) + "/config/octomap_multi_scan.yaml";

    YAML::Node node = YAML::LoadFile(yaml_path);
    res             = node["res"].as<double>();
    prob_hit        = node["prob_hit"].as<double>();
    prob_miss       = node["prob_miss"].as<double>();
    clamping_th_min = node["clamping_th_min"].as<double>();
    clamping_th_max = node["clamping_th_max"].as<double>();
    radius_filter_th= node["radius_filter_th"].as<double>(); 
    z_max_th        = node["z_max_th"].as<double>();
    distance_th     = node["distance_th"].as<double>();
    eps_angle       = node["eps_angle"].as<double>();
    do_ground_filtering = node["do_ground_filtering"].as<bool>();
    key_frame_gap   = node["key_frame_gap"].as<int>();
    p_ground_filter = std::make_shared<Ground_filter>(node);

    tree = new octomap::OcTree(res);
    tree->setClampingThresMax(clamping_th_max);
    tree->setClampingThresMin(clamping_th_min);
    tree->setProbHit(prob_hit);
    tree->setProbMiss(prob_miss);
    max_tree_depth = tree->getTreeDepth();

    grid_map.info.resolution = res;

    current_frame_pose = Eigen::Matrix4f::Identity();

    //initialize map saver
    std::string map_data_path = std::string(PROJECT_PATH) + "/data/grid_map.txt";
    std::string map_yaml_path = std::string(PROJECT_PATH) + "/data/grid_map.yaml";
    p_map_saver = std::make_shared<Map_saver>(map_yaml_path,map_data_path);
    // map_service_server = nh.advertiseService("mapping/save_map",&Gen_grid_map::save_map_callback,this);
}

void Gen_grid_map::ground_filtering(const pcl::PointCloud<pcl::PointXYZ>::Ptr& cloud_in){
     
     ground_cloud_in_map.reset(new pcl::PointCloud<pcl::PointXYZ>);
     non_ground_cloud_in_map.reset(new pcl::PointCloud<pcl::PointXYZ>);

     if(cloud_in->size() < 500){
        LOG(WARNING) << "cloud size too small , skip ground filtering.";
        non_ground_cloud_in_map = cloud_in;
        return;
    }
    else{
    pcl::ModelCoefficients::Ptr coefficients (new pcl::ModelCoefficients);
    pcl::PointIndices::Ptr inliers (new pcl::PointIndices);

    // Create the segmentation object and set up:
    pcl::SACSegmentation<pcl::PointXYZ> seg;
    seg.setOptimizeCoefficients (true);
    // TODO: maybe a filtering based on the surface normals might be more robust / accurate?
    seg.setModelType(pcl::SACMODEL_PERPENDICULAR_PLANE);
    seg.setMethodType(pcl::SAC_RANSAC);
    seg.setMaxIterations(200);
    seg.setDistanceThreshold (distance_th);//默认为0.04
    seg.setAxis(Eigen::Vector3f(0,0,1));
    seg.setEpsAngle(eps_angle);

    seg.setInputCloud(cloud_in);
    seg.segment(*inliers, *coefficients);

    pcl::ExtractIndices<pcl::PointXYZ> extract;
    extract.setInputCloud(cloud_in);
    extract.setIndices(inliers);
    extract.setNegative(false);
    extract.filter(*ground_cloud_in_map);

    extract.setNegative(true);
    extract.filter(*non_ground_cloud_in_map);

    //! 对non_ground在进行一次statistical outlier removal 去除外点
    pcl::StatisticalOutlierRemoval<pcl::PointXYZ> sor;
    sor.setInputCloud (non_ground_cloud_in_map);
    sor.setMeanK (10);
    sor.setStddevMulThresh (1.0);
    sor.filter(*non_ground_cloud_in_map);
}

}

void Gen_grid_map::DEBUG_pub_cloud_in_map(cloud_pose_msg::CloudPoseConstPtr msg_in){
    sensor_msgs::PointCloud2 non_ground_msg,ground_msg,full_msg;
    pcl::toROSMsg(*non_ground_cloud_in_map,non_ground_msg); 
    pcl::toROSMsg(*ground_cloud_in_map,ground_msg);
    //pcl::toROSMsg(*full_cloud_in_map,full_msg);
    non_ground_msg.header.frame_id =  "map";
    non_ground_msg.header.stamp    =  msg_in->header.stamp;
    ground_msg.header.frame_id     =  "map";
    ground_msg.header.stamp        =  msg_in->header.stamp;
    full_msg.header.frame_id       =  "map";
    full_msg.header.stamp          =  msg_in->header.stamp;

    //pub_full_cloud.      publish(full_msg);
    pub_ground_cloud.    publish(ground_msg);
    pub_non_ground_cloud.publish(non_ground_msg);
}

void Gen_grid_map::save_2d_map(){
    mtx.lock();
    //LOG(INFO) << "thread id in Gen_grid_map::save_2d_map() thread -> " << std::this_thread::get_id();


    p_map_saver->save_map(grid_map);
    //LOG(INFO) << "2D MAP SAVED.";

    mtx.unlock();
}

void Gen_grid_map::ray_tracing(){
    free_cells    .clear();
    occupied_cells.clear();
    
    octomap::KeyRay key_ray;
    for(int i = 0;i < non_ground_cloud_in_map->size();i++){
        octomap::point3d non_ground_point = octomap::point3d(non_ground_cloud_in_map->points[i].x,non_ground_cloud_in_map->points[i].y,non_ground_cloud_in_map->points[i].z);
        if(tree->computeRayKeys(sensor_origin,non_ground_point,key_ray)){
            free_cells.insert(key_ray.begin(),key_ray.end());
        }
        octomap::OcTreeKey end_key;
        if(tree->coordToKeyChecked(non_ground_point,end_key)){
            occupied_cells.insert(end_key);
        }
    }

    for(int i = 0;i < ground_cloud_in_map->size();i++){
        octomap::point3d ground_point = octomap::point3d(ground_cloud_in_map->points[i].x,ground_cloud_in_map->points[i].y,ground_cloud_in_map->points[i].z);
        if(tree->computeRayKeys(sensor_origin,ground_point,key_ray)){
            free_cells.insert(key_ray.begin(),key_ray.end());
        }
    }

    for(octomap::KeySet::iterator it = occupied_cells.begin() ; it != occupied_cells.end();it++){
        tree->updateNode(*it, true);
    }

    for(octomap::KeySet::iterator it = free_cells.begin() ; it != free_cells.end();it++){
        //对于本来没有被占据的栅格才更新为free
        //octomap::OcTreeNode* node = tree->search(*it);
        if(occupied_cells.find(*it) == occupied_cells.end()) 
        {
            tree->updateNode(*it, false);
        }
    }
    //tree->prune();
    tree->updateInnerOccupancy();
}

void Gen_grid_map::handle_occupied_node(octomap::OcTree::iterator& it){
    update_2d_map(it,true);
}   

void Gen_grid_map::handle_free_node(octomap::OcTree::iterator& it){
    update_2d_map(it,false);
}   

void Gen_grid_map::update_2d_map(const octomap::OcTree::iterator& it, bool occupied){
    // update 2D map (occupied always overrides):

  if (it.getDepth() == max_tree_depth){
    unsigned idx = mapIdx(it.getKey());
    if (occupied)
      grid_map.data[mapIdx(it.getKey())] = 100;
    else if (grid_map.data[idx] == -1){
      grid_map.data[idx] = 0;
    }

  } else{
    
    int intSize = 1 << (max_tree_depth - it.getDepth());
    //LOG(INFO) << "current_node depth = " << it.getDepth() << ".  int size =  " << intSize;

    octomap::OcTreeKey minKey=it.getIndexKey();
    for(int dx=0; dx < intSize; dx++){
      int i = (minKey[0]+dx - padded_min_key[0]);
      for(int dy=0; dy < intSize; dy++){
        unsigned idx = mapIdx(i, (minKey[1]+dy - padded_min_key[1]));//! 注意这个地方，padded_min_key or min_key??
        if (occupied)
          grid_map.data[idx] = 100;
        else if (grid_map.data[idx] == -1){
          grid_map.data[idx] = 0;
        }
      }
    }
  }
}

void Gen_grid_map::gen_2d_map(){


    double minX, minY, minZ, maxX, maxY, maxZ;
    tree->getMetricMax(maxX,maxY,maxZ);
    tree->getMetricMin(minX,minY,minZ);

    octomap::point3d minPt = octomap::point3d(minX,minY,minZ); //这个是以m为单位的
    octomap::point3d maxPt = octomap::point3d(maxX,maxY,maxZ);
    min_key = tree->coordToKey(minPt,max_tree_depth);
    max_key = tree->coordToKey(maxPt,max_tree_depth);

    
    minX = std::min(minX,0.);
    maxX = std::max(maxX,0.);
    minY = std::min(minY,0.);
    maxY = std::max(maxY,0.);
    minPt = octomap::point3d(minX, minY, minZ);
    maxPt = octomap::point3d(maxX, maxY, maxZ);

    octomap::OcTreeKey padded_max_key;
    if (!tree->coordToKeyChecked(minPt, max_tree_depth, padded_min_key)){//! m_paddedMinKey 代表了最边缘坐标所对应的key(上下限和0比较过)，注意！！！都是以最小leaf为单位来计算的
      ROS_ERROR("Could not create padded min OcTree key at %f %f %f", minPt.x(), minPt.y(), minPt.z());
      return;
    }
    if (!tree->coordToKeyChecked(maxPt, max_tree_depth, padded_max_key)){
      ROS_ERROR("Could not create padded max OcTree key at %f %f %f", maxPt.x(), maxPt.y(), maxPt.z());
      return;
    }

    
    assert(padded_max_key[0] >= max_key[0] && padded_max_key[1] >= max_key[1]);
    int mapOriginX = min_key[0] - padded_min_key[0];
    int mapOriginY = min_key[1] - padded_min_key[1];
    assert(mapOriginX >= 0 && mapOriginY >= 0);

    octomap::point3d origin = tree->keyToCoord(padded_min_key, max_tree_depth);
    


    grid_map.info.width = (padded_max_key[0] - padded_min_key[0]) + 1;
    grid_map.info.height = (padded_max_key[1] - padded_min_key[1]) + 1;
    grid_map.info.origin.position.x =  origin.x() - res*0.5;
    grid_map.info.origin.position.y =  origin.y() - res*0.5;
    //! *****************************************************************************************//
    grid_map.header.frame_id = "map";
    grid_map.header.stamp = ros::Time::now();
    // grid_map.info.width = (max_key[0] - min_key[0]) + 1;//! 这里有优化的空间 意思如果树因为剪枝 而深度减少，这个width就不是最小leaf的数量了。但这里暂不考虑
    // grid_map.info.height = (max_key[1] - min_key[1]) + 1;
    grid_map.info.resolution = res;
    // grid_map.info.origin.position.x = minPt.x();
    // grid_map.info.origin.position.y = minPt.y();
    ////grid_map.info.origin.position.z = 0;//! 注意 把地图原点设置在这个地方会不会有问题？？？？
    grid_map.data.clear();
    grid_map.data.resize(grid_map.info.width * grid_map.info.height,-1);

    for(octomap::OcTree::iterator it = tree->begin(max_tree_depth),end = tree->end();it != end;it++){
        octomap::OcTreeKey curr_key = octomap::OcTreeKey(it.getKey()[0] - padded_min_key[0],it.getKey()[1] - padded_min_key[1],0);
        
        if(tree->isNodeOccupied(*it)){
            handle_occupied_node(it);
        }
        else{
            handle_free_node(it);
        }
    }
}
void Gen_grid_map::filtering(const pcl::PointCloud<pcl::PointXYZ>::Ptr& _cloud_in,pcl::PointCloud<pcl::PointXYZ>::Ptr _cloud_out) {
  for(int i = 0;i < _cloud_in->size();i++){
        pcl::PointXYZ point(_cloud_in->points[i].x,_cloud_in->points[i].y,_cloud_in->points[i].z);
        double dis = point.x * point.x + point.y * point.y + point.z * point.z;
        if(dis < 25 & point.z < - 0.3) {
            continue;
        }
        if(dis < radius_filter_th * radius_filter_th && point.z < z_max_th && dis > 2){
            //! DEBUG
            //LOG(INFO) << "point.x/y/z = " << point.x << " " << point.y << " " << point.z; 
            _cloud_out->push_back(point);
            continue;
        }
    }
}
/*//! 流程
1、在lidar系下进行 地面滤波 ，得到ground cloud & non_ground_cloud

2、记录下sensor_origin 以及 将点云转换到map frame

3、对于地面点云、非地面点云进行ray_tracing

4、更新OcTree

5、根据OcTree重新生成OccupancyGrid

6、发布点云以及OccupancyGrid (in world frame)
*/
void Gen_grid_map::cloud_callback(const cloud_pose_msg::CloudPoseConstPtr& msg_in){
    frame_count++;
    if(frame_count % key_frame_gap != 0) {
        return;
    }
    TicToc tic;
    mtx.lock();
    if(msg_in->status != true) {
        LOG(INFO) << "ENU calibration not converged.skip this scan.";
        mtx.unlock();
        return;
    }
    LOG(INFO) << "cloud recieved.";
    //recieve data 
    sensor_origin = octomap::point3d(msg_in->pose.position.x,msg_in->pose.position.y,msg_in->pose.position.z);
    Eigen::Vector3f t_cloud_in(msg_in->pose.position.x,msg_in->pose.position.y,msg_in->pose.position.z);
    pcl::PointCloud<pcl::PointXYZ>::Ptr cloud_in(new pcl::PointCloud<pcl::PointXYZ>);
    pcl::fromROSMsg(msg_in->cloud,*cloud_in);
    Eigen::Quaternionf q_cloud_in(msg_in->pose.orientation.w,msg_in->pose.orientation.x,msg_in->pose.orientation.y,msg_in->pose.orientation.z);
    
    current_frame_pose.block<3,3>(0,0) = Eigen::Matrix3f(q_cloud_in);
    current_frame_pose.block<3,1>(0,3) = t_cloud_in;
    //pass throung filtering & ground filtering
    std::vector<int> indices;
    pcl::PointCloud<pcl::PointXYZ>::Ptr filtered_cloud(new pcl::PointCloud<pcl::PointXYZ>);
    filtering(cloud_in,filtered_cloud);

    pcl::VoxelGrid<pcl::PointXYZ> sor;
    sor.setLeafSize(0.1,0.1,0.1);
    sor.setInputCloud(filtered_cloud);
    sor.filter(*filtered_cloud);

    if(filtered_cloud->size() < 100){
        LOG(WARNING) << "cloud size too small after filtering.ERROR!";
        return;
    }
    cloud_in = filtered_cloud;

    //GROUND FILTERING & transfomr cloud into map frame
    //ground_filtering(cloud_in) 
    //通过配置文件对是否进行地面滤波进行设置
    if(do_ground_filtering) {
        p_ground_filter->ground_filtering(cloud_in,ground_cloud_in_map,non_ground_cloud_in_map);
        pcl::transformPointCloud(*ground_cloud_in_map,*ground_cloud_in_map,current_frame_pose);
        pcl::transformPointCloud(*non_ground_cloud_in_map,*non_ground_cloud_in_map,current_frame_pose);
    }
    else {
        non_ground_cloud_in_map = cloud_in;
        ground_cloud_in_map.reset(new pcl::PointCloud<pcl::PointXYZ>);
        pcl::transformPointCloud(*non_ground_cloud_in_map,*non_ground_cloud_in_map,current_frame_pose);
    }
    // p_ground_filter->ground_filtering(cloud_in,ground_cloud_in_map,non_ground_cloud_in_map);
    // pcl::transformPointCloud(*ground_cloud_in_map,*ground_cloud_in_map,current_frame_pose);
    // pcl::transformPointCloud(*non_ground_cloud_in_map,*non_ground_cloud_in_map,current_frame_pose);
    full_cloud_in_map.reset(new pcl::PointCloud<pcl::PointXYZ>(*cloud_in));
    pcl::transformPointCloud(*full_cloud_in_map,*full_cloud_in_map,current_frame_pose);

    DEBUG_pub_cloud_in_map(msg_in);
    //RAY TRACING
    ray_tracing();
    //generate a new (not incremental) 2d map from octree
    gen_2d_map();
    //publish grid map
    pub_grid_map.publish(grid_map);
    mtx.unlock();
    LOG(INFO) << "process cloud call-back time used : " << tic.toc();
}





